from fastmcp import FastMCP
from fastmcp.prompts.prompt import PromptMessage
from fastmcp.server.auth import AccessToken
from fastmcp.server.auth.providers.bearer import RSAKeyPair, BearerAuthProvider
from fastmcp.server.dependencies import get_access_token
from zhipuai import ZhipuAI
from agent.env_utils import ZHIPU_API_KEY


# 1.生成RSA密钥对
key_pair = RSAKeyPair.generate()
print(key_pair.public_key)

# 2.配置提供认证方
auth = BearerAuthProvider(
    public_key=key_pair.public_key,     # 公钥用于校验签名
    issuer='https://www.laoxiao.com',       # 令牌签发方标识
    audience='my-dev-server'        # 服务商的一个标识
)

# 3.服务器模拟生成一个token
token = key_pair.create_token(
    subject='dev-user',
    issuer='https://www.laoxiao.com',
    audience='my-dev-server',
    scopes=['laoxiao', 'invoke_tools'],
    expires_in_seconds=3600,
)
print(f"Test token: {token}")


zhipuai_client = ZhipuAI(api_key=ZHIPU_API_KEY)
server = FastMCP(
    name='my_py_mcp_2',
    instructions='我的Python代码实现的MCP服务器',
    auth=auth,      # 服务用于校验token的
)


@server.tool(name='zhipuai_search')
def my_search(query: str):
    """搜索互联网上的内容,包括实时天气等"""
    try:
        print("执行我的Python中的工具，输入的参数为:", query)

        # 得到了验证通过之后的
        access_token: AccessToken = get_access_token()
        if access_token:
            print(f"AccessToken: {access_token}")
            print(access_token.scopes)
        else:
            return '由于没有权限，所以不能搜索到任何内容！，请客户端传入有效token'
        response = zhipuai_client.web_search.web_search(
            search_engine="search_pro",
            search_query=query
        )
        # print(response)
        if response.search_result:
            return "\n\n".join([d.content for d in response.search_result])
        return '没有搜索到任何内容！'
    except Exception as e:
        print(e)
        return '没有搜索到任何内容！'


@server.tool()
def say_hello(username: str) -> str:
    """给指定用户打个招呼"""
    return f"{username}， 你好，今天天气不错！"


@server.prompt
def ask_about_topic(topic: str):
    """生成请求解释特定主题的用户消息模板"""
    return f"能否请您解释一下'{topic}',这个概念？"


# 高级提示模板：返回结构化消息对象
@server.prompt
def generate_code_request(language: str, task_description: str):
    """生成代码编写请求的用户消息模板"""
    content = f"请用{language}编写一个实现以下功能的函数：{task_description}"
    return PromptMessage(
        role='user',
        content=content
    )


@server.resource("resource://config")
def get_config() -> dict:
    """以JSON格式返回应用配置"""
    return {
        "theme": "dark",  # 界面主题配置
        "version": "1.2.0",  # 当前版本号
        "features": ["tools", "resources"],  # 已启用的功能模块
    }
